package com.wechat.backend.ws;

import com.fasterxml.jackson.databind.ObjectMapper;
import com.wechat.backend.service.MessageService;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.*;
import org.springframework.web.socket.handler.TextWebSocketHandler;

import java.util.*;
import java.util.concurrent.ConcurrentHashMap;

@Component
public class ChatWebSocketHandler extends TextWebSocketHandler {

    // 存储每个用户的 WebSocketSession
    private static final Map<Integer, WebSocketSession> userSessionMap = new ConcurrentHashMap<>();

    @Autowired
    private MessageService messageService; // 只用，不改

    private final ObjectMapper objectMapper = new ObjectMapper();

    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        // 获取 userId
        String path = session.getUri().getPath();
        String userIdStr = path.substring(path.lastIndexOf('/') + 1);
        Integer userId = Integer.valueOf(userIdStr);
        userSessionMap.put(userId, session);
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
        userSessionMap.values().removeIf(s -> s.equals(session));
    }

    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
        // 解析消息
        Map<String, Object> msgMap = objectMapper.readValue(message.getPayload(), Map.class);
        Integer fromId = (Integer) msgMap.get("fromId");
        Integer toId = (Integer) msgMap.get("toId");
        String content = (String) msgMap.get("content");

        // 只调用已有的消息保存方法，不改动原有实现
        messageService.saveMessage(fromId, toId, content);

        Map<String, Object> resp = new HashMap<>();
        resp.put("fromId", fromId);
        resp.put("toId", toId);
        resp.put("content", content);
        resp.put("createTime", new Date());

        String respJson = objectMapper.writeValueAsString(resp);

        WebSocketSession fromSession = userSessionMap.get(fromId);
        if (fromSession != null && fromSession.isOpen()) {
            fromSession.sendMessage(new TextMessage(respJson));
        }
        WebSocketSession toSession = userSessionMap.get(toId);
        if (toSession != null && toSession.isOpen() && !Objects.equals(fromId, toId)) {
            toSession.sendMessage(new TextMessage(respJson));
        }
    }
} 